'''添加格网有效数据个数信息'''
import numpy as np
import netCDF4 as nc
import os
import glob
from tqdm import tqdm

# 配置参数
INPUT_DIR = "/mnt/datastore/liudddata/season/3_5"
# INPUT_DIR = "/mnt/datastore/liudddata/result/20200101test"
OUTPUT_NC = "/mnt/datastore/liudddata/season/average_fycth_3_5new.nc"
# OUTPUT_NC = "/mnt/datastore/liudddata/season/average_fycbh_20200101new.nc"
# INPUT_DIR = "/mnt/datastore/liudddata/result/20190212_01new"
# OUTPUT_NC = "/mnt/datastore/liudddata/season/average_cbh_annual_corrected.nc"
VARIABLE_NAME = "predicted"
# VARIABLE_NAME = "cth"
FILL_VALUE = -999.9


def process_2d_grid():
    # 获取文件列表
    file_list = sorted(glob.glob(os.path.join(INPUT_DIR, "*_predicted_2d.nc")))
    if not file_list:
        raise ValueError("未找到输入文件")

    # 读取第一个文件获取网格信息
    with nc.Dataset(file_list[0], 'r') as template_ds:
        # 验证维度结构
        if 'y' not in template_ds.dimensions or 'x' not in template_ds.dimensions:
            raise RuntimeError("原始文件维度不符合预期结构")

        # 获取维度大小
        y_dim = len(template_ds.dimensions['y'])
        x_dim = len(template_ds.dimensions['x'])

        # 关键修改：正确读取二维坐标并保持维度顺序
        lat_2d = template_ds['lat'][:, :]  # 显式读取二维数组
        lon_2d = template_ds['lon'][:, :]

        # 验证坐标维度顺序
        if lat_2d.shape != (y_dim, x_dim):
            raise RuntimeError("纬度数据维度顺序异常")

    # 初始化累加数组
    sum_data = np.zeros((y_dim, x_dim), dtype=np.float64)
    count_data = np.zeros((y_dim, x_dim), dtype=np.uint32)

    # 逐文件处理
    for file_path in tqdm(file_list, desc="处理进度"):
        with nc.Dataset(file_path, 'r') as ds:
            # 验证数据维度
            if ds[VARIABLE_NAME].shape != (y_dim, x_dim):
                continue

            # 读取数据并处理
            data = ds[VARIABLE_NAME][:, :].astype(np.float32)  # 显式二维索引
            # 添加云底高度小于 0 为无效值的判断
            valid_mask = (~np.isclose(data, FILL_VALUE)) & (~np.isnan(data)) & (data > 0)

            # 累加计算
            sum_data += np.where(valid_mask, data, 0)
            count_data += valid_mask.astype(np.uint32)

    # 计算平均值
    with np.errstate(divide='ignore', invalid='ignore'):
        mean_data = np.divide(sum_data, count_data, where=count_data > 0)
    mean_data = np.where(count_data > 0, mean_data, FILL_VALUE).astype(np.float32)

    # 创建输出文件
    with nc.Dataset(OUTPUT_NC, 'w', format='NETCDF4') as out_ds:
        # 定义维度（保持与输入文件相同顺序）
        out_ds.createDimension('y', y_dim)
        out_ds.createDimension('x', x_dim)

        # 创建坐标变量（显式指定维度顺序）
        lat_var = out_ds.createVariable('lat', 'f4', ('y', 'x'), zlib=True)
        lon_var = out_ds.createVariable('lon', 'f4', ('y', 'x'), zlib=True)

        # 二维数组直接赋值
        lat_var[:, :] = lat_2d
        lon_var[:, :] = lon_2d

        # 添加坐标属性
        lat_var.units = "degrees_north"
        lat_var.long_name = "latitude"
        lon_var.units = "degrees_east"
        lon_var.long_name = "longitude"

        # 创建数据变量
        cbh_var = out_ds.createVariable(
            'mean_cbh', 'f4', ('y', 'x'),
            fill_value=FILL_VALUE,
            zlib=True,
            complevel=2
        )
        cbh_var[:, :] = mean_data  # 二维赋值

        # 创建记录有效数据个数的变量
        valid_count_var = out_ds.createVariable(
            'valid_count', 'u4', ('y', 'x'),  # 'u4' 表示无符号32位整数
            zlib=True,
            complevel=2
        )
        valid_count_var[:, :] = count_data  # 将记录的有效数据个数写入变量

        # 添加全局属性
        out_ds.title = "FY-4A网格年平均云底高度"
        out_ds.grid_type = "不规则二维网格"
        out_ds.source_files = ";".join([os.path.basename(f) for f in file_list])


if __name__ == "__main__":
    try:
        process_2d_grid()
        print(f"处理完成，结果保存至：{OUTPUT_NC}")
    except Exception as e:
        print(f"处理失败：{str(e)}")
